# 演示Unitary Transform
import argparse
from typing import Dict
import matplotlib.pyplot as plt
import numpy as np
from mpl_toolkits.mplot3d import Axes3D

class Chp01Sec03(object):
    def __init__(self):
        self.name = ''

    @staticmethod
    def startup(params:Dict = {}) -> None:
        print(f'Unitary Transformation Demo v0.0.1')
        plt.rcParams['figure.figsize'] = [16, 8]
        plt.rcParams.update({'font.size': 18})
        # 定义旋转角度
        theta = np.array([np.pi/15, -np.pi/9, -np.pi/20])
        # 定义放缩系数：x轴放大3倍，y轴不变，z轴缩小为0.5倍
        Sigma = np.diag([3, 1, 0.5]) # scale x, then y, then z
        # Rotation about x axis
        Rx = np.array([[1, 0, 0],
                    [0, np.cos(theta[0]), -np.sin(theta[0])],
                    [0, np.sin(theta[0]), np.cos(theta[0])]])
        # Rotation about y axis
        Ry = np.array([[np.cos(theta[1]), 0, np.sin(theta[1])],
                    [0, 1, 0],
                    [-np.sin(theta[1]), 0, np.cos(theta[1])]])
        # Rotation about z axis
        Rz = np.array([[np.cos(theta[2]), -np.sin(theta[2]), 0],
                    [np.sin(theta[2]), np.cos(theta[2]), 0],
                    [0, 0, 1]])
        # Rotate and scale
        X = Rz @ Ry @ Rx @ Sigma # \hat{U}\hat{\Sigma}
        # Plot sphere 定义并绘制单位立方球体
        fig = plt.figure()
        ax1 = fig.add_subplot(121, projection='3d')
        u = np.linspace(-np.pi, np.pi, 100)
        v = np.linspace(0, np.pi, 100)
        x = np.outer(np.cos(u), np.sin(v))
        y = np.outer(np.sin(u), np.sin(v))
        z = np.outer(np.ones(np.size(u)), np.cos(v))
        # Plot the surface
        surf1 = ax1.plot_surface(x, y, z, cmap='jet',alpha=0.6,facecolors=plt.cm.jet(z),linewidth=0.5,rcount=30,ccount=30)
        surf1.set_edgecolor('k')
        ax1.set_xlim3d(-2, 2)
        ax1.set_ylim3d(-2, 2)
        ax1.set_zlim3d(-2, 2)
        # Define the matrix X = [xR, yR, zR]
        xR = np.zeros_like(x)
        yR = np.zeros_like(y)
        zR = np.zeros_like(z)
        # X是Unitary转换，vec为待转换单位向量
        for i in range(x.shape[0]):
            for j in range(x.shape[1]):
                vec = [x[i,j], y[i,j], z[i,j]]
                vecR = X @ vec
                xR[i,j] = vecR[0]
                yR[i,j] = vecR[1]
                zR[i,j] = vecR[2]
        # 绘制转换后的图像
        ax2 = fig.add_subplot(122, projection='3d')
        surf2 = ax2.plot_surface(xR, yR, zR, cmap='jet',alpha=0.6,linewidth=0.5,facecolors=plt.cm.jet(z),rcount=30,ccount=30)
        surf2.set_edgecolor('k')
        ax2.set_xlim3d(-2, 2)
        ax2.set_ylim3d(-2, 2)
        ax2.set_zlim3d(-2, 2)
        plt.show()

def main(params:Dict = {}) -> None:
    Chp01Sec03.startup(params=params)

def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--run_mode', action='store',
        type=int, default=1, dest='run_mode',
        help='run mode'
    )
    return parser.parse_args()

if '__main__' == __name__:
    args = parse_args()
    params = vars(args)
    main(params=params)